"""
Author : sunshicheng
DateTime : 2021/6/1.4:25 下午
FileName : grpc_locust.py
Desc: 
"""

import logging
import time
import grpc
import helloworld_pb2
import helloworld_pb2_grpc
import json
from locust import (TaskSet, task, events, Locust)
from gevent._semaphore import Semaphore
import random

log_fmt = "[%(levelname)s]%(asctime)s line %(lineno)d :\n%(message)s"
c_fmt = "[%(levelname)s]%(asctime)s %(filename)s.%(funcName)s():line %(lineno)d :\n%(message)s"
date_format = "%Y-%m-%d %H:%M:%S %a"
# 设置控制台输出level
logging.basicConfig(level=logging.INFO,
                    format=c_fmt,
                    datefmt=date_format
                    )
all_locusts_spawned = Semaphore()
all_locusts_spawned.acquire()
host = "192.168.110.135"
port = "50051"


def on_hatch_complete(**kwargs):
    all_locusts_spawned.release()


events.hatch_complete += on_hatch_complete


def run():
    """test grpc server demo """
    with grpc.insecure_channel('192.168.110.135:50051') as channel:
        stub = helloworld_pb2_grpc.GreeterStub(channel)
        req = {"name": "jack", "id": 1001}
        body = json.dumps(req).encode("utf-8")
        response = stub.SayHello(helloworld_pb2.HelloRequest(name=body))
    print(response.message)  # str


class GrpcClient(object):
    """overide client"""

    def __init__(self):
        self.host = host
        self.port = port

    def grpc_request(self, body):
        start_time = int(time.time())
        response = None
        try:
            address = "{}:{}".format(host, port)
            channel = grpc.insecure_channel(address)
            # p2_grpc.Stub class implement new locust client
            client_stub = helloworld_pb2_grpc.GreeterStub(channel=channel)
            request_object = helloworld_pb2.HelloRequest(name=body)
            back_req = client_stub.SayHello(request_object)
            response = back_req.message
            elapsed = int((time.time() - start_time) * 1000)
            text = json.loads(response)
            logging.info("get response is {}".format(response))
            if text["retcode"] != 0:
                raise Exception("response get not expect,actual is {}".format(text))
            events.request_success.fire(
                request_type='grpc',
                name=r'/SayHello',
                response_time=elapsed,
                response_length=0
            )
        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            events.request_failure.fire(
                request_type='grpc',
                name='/SayHello',
                response_time=total_time,
                exception=e
            )

        return response


class GrpcLocust(Locust):
    """ overide Locust to implement GrpcLocust"""

    def __init__(self, *args, **kwargs):
        super(GrpcLocust, self).__init__()
        self.client = GrpcClient()


class GrpcUserBehavior(TaskSet):
    """ super TaskSet class implement new GrpcTaskSet"""

    def on_task(self):
        """wait task event spawn"""
        all_locusts_spawned.wait()

    def on_stop(self):
        pass

    @task
    def inference_task(self):
        code = random.choice([101, 102, 0, 0])
        data = {"name": "{}".format("ZhangSan"), "retcode": code}
        body = json.dumps(data).encode("utf-8")
        response_msg = self.client.grpc_request(body)
        # print(response_msg)


class WebsiteUser(GrpcLocust):
    task_set = GrpcUserBehavior
    min_wait = 200  # think time ms
    max_wait = 500
